#coding:utf-8

import numpy as np
import pandas as pd

import sys
np.set_printoptions(threshold=sys.maxsize)

from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error

df_train=pd.read_csv("./dataset/train.csv")
df_test=pd.read_csv("./dataset/test.csv")

df_train["Date"]=pd.to_datetime(df_train["Date"])
df_test.Date=pd.to_datetime(df_test["Date"])

df_test.drop('Province/State', axis = 1, inplace = True)
date_filter = df_train['Date'] < df_test['Date'].min()
df_train = df_train.loc[date_filter]
train_country_date = df_train.groupby(['Country/Region', 'Date', 'Lat', 'Long'], as_index=False)['ConfirmedCases', 'Fatalities'].sum()

train_country_date['Month'] = train_country_date['Date'].dt.month
train_country_date['Day'] = train_country_date['Date'].dt.day
train_country_date['Day_Week'] = train_country_date['Date'].dt.dayofweek
train_country_date['quarter'] = train_country_date['Date'].dt.quarter
train_country_date['dayofyear'] = train_country_date['Date'].dt.dayofyear
train_country_date['weekofyear'] = train_country_date['Date'].dt.weekofyear
df_test['Month'] = df_test['Date'].dt.month
df_test['Day'] = df_test['Date'].dt.day
df_test['Day_Week'] = df_test['Date'].dt.dayofweek
df_test['quarter'] = df_test['Date'].dt.quarter
df_test['dayofyear'] = df_test['Date'].dt.dayofyear
df_test['weekofyear'] = df_test['Date'].dt.weekofyear

labels = ['Country/Region', 'Lat', 'Long', 'Date', 'Month', 'Day', 'Day_Week','quarter', 'dayofyear', 'weekofyear']
df_train_clean = train_country_date[labels]
df_test_clean = df_test[labels]
data_clean = pd.concat([df_train_clean, df_test_clean], axis = 0)

enc = LabelEncoder()
data_clean['Country'] = enc.fit_transform(data_clean['Country/Region'])

data_clean.drop(['Country/Region', 'Date'], axis = 1, inplace=True)

index_split = df_train.shape[0]
data_train_clean = data_clean[:index_split]
data_test_clean = data_clean[index_split:]

x = data_train_clean[['Lat', 'Long', 'Month', 'Day', 'Day_Week','quarter', 'dayofyear', 'weekofyear', 'Country']]
y_case = df_train['ConfirmedCases']
y_fatal = df_train['Fatalities']

x_train, x_test, y_train, y_test = train_test_split(x, y_case, test_size = 0.3, random_state = 42)

x_train_fatal, x_test_fatal, y_train_fatal, y_test_fatal = train_test_split(x, y_fatal, test_size = 0.3, random_state = 42)

rf = RandomForestRegressor(n_estimators =100)
rf.fit(x_train, y_train.values)

rf.score(x_train, y_train)
rf.score(x_test, y_test)

y_pred_train = rf.predict(x_train)
print(mean_squared_error(y_train, y_pred_train))

rf.fit(x, y_case.values.ravel())

#############用最终模型去预测
rf_pred_case = rf.predict(data_test_clean)
rf.fit(x, y_fatal.values.ravel())
rf_pred_fatal = rf.predict(data_test_clean)

randomforest=pd.read_csv("test.csv")
randomforest['ConfirmedCases'] = rf_pred_case
randomforest['Fatalities'] = rf_pred_fatal
randomforest.to_csv('randomforest.csv', index = False)
